{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "05aebd11-2df1-4daf-a598-8359b2d92c64",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pylint: disable=maybe-no-member\n",
    "# from plot_model import plot_results\n",
    "import torch\n",
    "import torchvision\n",
    "import torch.nn as nn\n",
    "import matplotlib.pyplot as plt\n",
    "from torchvision import datasets\n",
    "# from torchvision import transforms\n",
    "import torchvision.transforms as transforms\n",
    "from torch.autograd import Variable\n",
    "import time\n",
    "# from sklearn.preprocessing import LabelEncoder, OneHotEncoder\n",
    "from sklearn.metrics import accuracy_score, recall_score, roc_curve\n",
    "from astropy.io import fits\n",
    "import numpy as np\n",
    "import glob\n",
    "import torch.utils.data as Data\n",
    "from sklearn.model_selection import train_test_split\n",
    "from glob import glob\n",
    "import warnings\n",
    "import itertools\n",
    "import logging\n",
    "import time # 时间\n",
    "import os # 路径\n",
    "from sklearn.metrics import confusion_matrix ,classification_report\n",
    "from sklearn.preprocessing import LabelEncoder, OneHotEncoder\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "63a53105-a2e8-4e4d-aa51-59aa0774751d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_data(file_all=None, MK=True):\n",
    "    flux = np.load('flux.npy')\n",
    "    scls = np.load('spectypes.npy')\n",
    "\n",
    "    #print(flux.shape, scls.shape)\n",
    "\n",
    "    fluxTR, fluxTE, clsTR, clsTE = train_test_split(flux, scls, test_size=0.2)\n",
    "    Xtrain1 = torch.from_numpy(flux)  # numpy 转成 torch 类型\n",
    "    Xtest1 = torch.from_numpy(fluxTE)\n",
    "    ytrain1 = torch.from_numpy(scls)\n",
    "    ytest1 = torch.from_numpy(clsTE)\n",
    "    # transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean=[0.5,],std=[0.5,])])\n",
    "    torch_dataset_train = Data.TensorDataset(Xtrain1, ytrain1)\n",
    "    torch_dataset_test = Data.TensorDataset(Xtest1, ytest1)\n",
    "\n",
    "    data_loader_train = torch.utils.data.DataLoader(dataset=torch_dataset_train, batch_size=batch_size, shuffle=True)\n",
    "    data_loader_test = torch.utils.data.DataLoader(dataset=torch_dataset_test, batch_size=batch_size, shuffle=True)\n",
    "    return data_loader_train, data_loader_test, clsTR.shape[0], batch_size\n",
    "\n",
    "\n",
    "#     return onehot_encoded\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a89175f0-8dc5-4919-a18b-f13292f57024",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def get_variable(x):\n",
    "    x = Variable(x)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2c5da470-2249-4645-a51f-227687b11f7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RNN_Model(torch.nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, layer_size, num_class):\n",
    "        super(RNN_Model, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        self.layer_size = layer_size\n",
    "        self.rnn = torch.nn.RNN(input_size,hidden_size,layer_size,batch_first=True)\n",
    "        self.dense = torch.nn.Sequential(\n",
    "            torch.nn.ReLU(),\n",
    "            torch.nn.Linear( hidden_size, num_class),\n",
    "            torch.nn.Softmax()\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        # print(x.shape)\n",
    "        # 初始化隐藏状态\n",
    "        h0 = torch.zeros(self.layer_size, x.size(0), self.hidden_size).requires_grad_()\n",
    "        # 前向传播RNN层\n",
    "        out, _ = self.rnn(x, h0.detach())\n",
    "        # 取最后一个时间步的输出\n",
    "        out = out[:, -1, :]\n",
    "        # 前向传播全连接层\n",
    "        out = self.dense(out)\n",
    "        return out\n",
    "        #x = x.view(-1, *x.shape[1:])\n",
    "        #x, _ = self.rnn(x)\n",
    "        # print(x.shape)    \n",
    "        #x = x.contiguous().view(-1, *x.shape[1:])\n",
    "        #x = self.dense(x)\n",
    "        #return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6eed62fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_RNN_module(device, num_class, num_epochs, batch_size, learning_rate, train_loader, test_loader, clsTR, clsTE):\n",
    "    # 初始化模型、损失函数和优化器\n",
    "    input_size = 7781  # 输入特征维度\n",
    "    hidden_size = 128  # 隐藏层维度\n",
    "    num_layers = 2     # RNN 层数\n",
    "    model = RNN_Model(input_size, hidden_size, num_layers, num_class).to(device)\n",
    "    loss_func = nn.CrossEntropyLoss()\n",
    "    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9)\n",
    "\n",
    "    # 训练循环\n",
    "    for epoch in range(num_epochs):\n",
    "        TR_loss=0.0\n",
    "        TR_correct=0\n",
    "        TR_total = 0.0\n",
    "\n",
    "        training_time = time.time()\n",
    "        for i, (data) in enumerate(train_loader):\n",
    "            optimizer.zero_grad()\n",
    "            X_train, y_train = data\n",
    "            X_train, y_train = X_train.to(device), y_train.to(device)\n",
    "\n",
    "            # 检查并调整输入张量的形状\n",
    "            if len(X_train.shape) == 2:\n",
    "                X_train = X_train.unsqueeze(1)  # 添加一个维度以匹配 RNN 输入要求\n",
    "\n",
    "            outputs = model(X_train)\n",
    "            loss = loss_func(outputs, y_train.long())\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            TR_loss += loss.item()\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            TR_total += y_train.size(0)\n",
    "            TR_correct += (predicted == y_train.long()).sum().item()\n",
    "\n",
    "            if (i + 1) % 100 == 0:\n",
    "                print(f'Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}], Loss: {loss.item():.4f}')\n",
    "\n",
    "        # 测试模型\n",
    "        model.eval()\n",
    "        with torch.no_grad():\n",
    "            TE_loss = 0.0\n",
    "            TE_correct=0\n",
    "            TE_total = 0\n",
    "            for data in test_loader:\n",
    "                X_test, y_test = data\n",
    "                X_test, y_test = X_test.to(device), y_test.to(device)\n",
    "\n",
    "                # 检查并调整输入张量的形状\n",
    "                if len(X_test.shape) == 2:\n",
    "                    X_test = X_test.unsqueeze(1)  # 添加一个维度以匹配 RNN 输入要求\n",
    "\n",
    "                outputs = model(X_test)\n",
    "                loss = loss_func(outputs, y_test.long())\n",
    "                TE_loss += loss.item()\n",
    "                _, predicted = torch.max(outputs.data, 1)\n",
    "                TE_total += y_test.size(0)\n",
    "                TE_correct += (predicted == y_test.long()).sum().item()\n",
    "\n",
    "            print(f\"Epoch {epoch+1}:\\n\"\n",
    "                  f\"Train Loss: {TR_loss:.4f}, train Acc: {100 * TR_correct / TR_total:.2f}%\\n\"\n",
    "                  f\"test Loss: {TE_loss:.4f}, test Acc: {100 * TE_correct / TE_total:.2f}%\")\n",
    "            print(\"-\" * 50)\n",
    "\n",
    "        model.train()\n",
    "\n",
    "    print(f\"Training time: {time.time() - training_time} seconds\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "76a75aeb-1149-407a-a930-50d802d00635",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1:\n",
      "Train Loss: 100.5780, train Acc: 48.83%\n",
      "test Loss: 19.7554, test Acc: 56.83%\n",
      "--------------------------------------------------\n",
      "Epoch 2:\n",
      "Train Loss: 95.5357, train Acc: 56.45%\n",
      "test Loss: 18.9878, test Acc: 58.92%\n",
      "--------------------------------------------------\n",
      "Epoch 3:\n",
      "Train Loss: 92.2801, train Acc: 59.60%\n",
      "test Loss: 18.4172, test Acc: 61.08%\n",
      "--------------------------------------------------\n",
      "Epoch 4:\n",
      "Train Loss: 89.8721, train Acc: 61.63%\n",
      "test Loss: 18.0008, test Acc: 63.58%\n",
      "--------------------------------------------------\n",
      "Epoch 5:\n",
      "Train Loss: 87.9544, train Acc: 63.87%\n",
      "test Loss: 17.6593, test Acc: 65.33%\n",
      "--------------------------------------------------\n",
      "Epoch 6:\n",
      "Train Loss: 86.3024, train Acc: 65.58%\n",
      "test Loss: 17.3258, test Acc: 66.67%\n",
      "--------------------------------------------------\n",
      "Epoch 7:\n",
      "Train Loss: 84.9507, train Acc: 68.33%\n",
      "test Loss: 17.0383, test Acc: 69.00%\n",
      "--------------------------------------------------\n",
      "Epoch 8:\n",
      "Train Loss: 83.6668, train Acc: 70.20%\n",
      "test Loss: 16.7765, test Acc: 70.33%\n",
      "--------------------------------------------------\n",
      "Epoch 9:\n",
      "Train Loss: 82.5248, train Acc: 72.17%\n",
      "test Loss: 16.5452, test Acc: 72.92%\n",
      "--------------------------------------------------\n",
      "Epoch 10:\n",
      "Train Loss: 81.5456, train Acc: 73.10%\n",
      "test Loss: 16.3361, test Acc: 74.33%\n",
      "--------------------------------------------------\n",
      "Epoch 11:\n",
      "Train Loss: 80.5600, train Acc: 74.78%\n",
      "test Loss: 16.1601, test Acc: 76.08%\n",
      "--------------------------------------------------\n",
      "Epoch 12:\n",
      "Train Loss: 79.7234, train Acc: 75.85%\n",
      "test Loss: 15.9655, test Acc: 77.08%\n",
      "--------------------------------------------------\n",
      "Epoch 13:\n",
      "Train Loss: 78.8548, train Acc: 76.33%\n",
      "test Loss: 15.7899, test Acc: 78.33%\n",
      "--------------------------------------------------\n",
      "Epoch 14:\n",
      "Train Loss: 78.0791, train Acc: 77.13%\n",
      "test Loss: 15.6118, test Acc: 78.92%\n",
      "--------------------------------------------------\n",
      "Epoch 15:\n",
      "Train Loss: 77.2736, train Acc: 77.62%\n",
      "test Loss: 15.5118, test Acc: 78.75%\n",
      "--------------------------------------------------\n",
      "Epoch 16:\n",
      "Train Loss: 76.8714, train Acc: 78.07%\n",
      "test Loss: 15.3775, test Acc: 79.67%\n",
      "--------------------------------------------------\n",
      "Epoch 17:\n",
      "Train Loss: 76.1400, train Acc: 78.27%\n",
      "test Loss: 15.2366, test Acc: 80.25%\n",
      "--------------------------------------------------\n",
      "Epoch 18:\n",
      "Train Loss: 75.6281, train Acc: 78.72%\n",
      "test Loss: 15.1661, test Acc: 80.42%\n",
      "--------------------------------------------------\n",
      "Epoch 19:\n",
      "Train Loss: 75.0709, train Acc: 78.97%\n",
      "test Loss: 15.0303, test Acc: 80.08%\n",
      "--------------------------------------------------\n",
      "Epoch 20:\n",
      "Train Loss: 74.5801, train Acc: 79.38%\n",
      "test Loss: 14.9184, test Acc: 80.33%\n",
      "--------------------------------------------------\n",
      "Epoch 21:\n",
      "Train Loss: 73.9743, train Acc: 79.87%\n",
      "test Loss: 14.7571, test Acc: 81.17%\n",
      "--------------------------------------------------\n",
      "Epoch 22:\n",
      "Train Loss: 73.3046, train Acc: 80.48%\n",
      "test Loss: 14.6605, test Acc: 80.92%\n",
      "--------------------------------------------------\n",
      "Epoch 23:\n",
      "Train Loss: 72.8849, train Acc: 80.77%\n",
      "test Loss: 14.5810, test Acc: 81.67%\n",
      "--------------------------------------------------\n",
      "Epoch 24:\n",
      "Train Loss: 72.4099, train Acc: 81.37%\n",
      "test Loss: 14.4602, test Acc: 82.50%\n",
      "--------------------------------------------------\n",
      "Epoch 25:\n",
      "Train Loss: 72.0729, train Acc: 81.63%\n",
      "test Loss: 14.3964, test Acc: 82.58%\n",
      "--------------------------------------------------\n",
      "Epoch 26:\n",
      "Train Loss: 71.6683, train Acc: 81.85%\n",
      "test Loss: 14.3379, test Acc: 82.75%\n",
      "--------------------------------------------------\n",
      "Epoch 27:\n",
      "Train Loss: 71.2173, train Acc: 82.32%\n",
      "test Loss: 14.2467, test Acc: 83.33%\n",
      "--------------------------------------------------\n",
      "Epoch 28:\n",
      "Train Loss: 70.8763, train Acc: 82.67%\n",
      "test Loss: 14.1604, test Acc: 83.42%\n",
      "--------------------------------------------------\n",
      "Epoch 29:\n",
      "Train Loss: 70.5759, train Acc: 83.03%\n",
      "test Loss: 14.1195, test Acc: 83.83%\n",
      "--------------------------------------------------\n",
      "Epoch 30:\n",
      "Train Loss: 70.1839, train Acc: 83.30%\n",
      "test Loss: 14.0306, test Acc: 84.50%\n",
      "--------------------------------------------------\n",
      "Epoch 31:\n",
      "Train Loss: 69.8474, train Acc: 83.52%\n",
      "test Loss: 13.9891, test Acc: 84.58%\n",
      "--------------------------------------------------\n",
      "Epoch 32:\n",
      "Train Loss: 69.5454, train Acc: 83.83%\n",
      "test Loss: 13.9244, test Acc: 84.92%\n",
      "--------------------------------------------------\n",
      "Epoch 33:\n",
      "Train Loss: 69.1479, train Acc: 84.08%\n",
      "test Loss: 13.7961, test Acc: 85.33%\n",
      "--------------------------------------------------\n",
      "Epoch 34:\n",
      "Train Loss: 68.7543, train Acc: 84.57%\n",
      "test Loss: 13.7441, test Acc: 86.00%\n",
      "--------------------------------------------------\n",
      "Epoch 35:\n",
      "Train Loss: 68.4263, train Acc: 84.80%\n",
      "test Loss: 13.7078, test Acc: 85.75%\n",
      "--------------------------------------------------\n",
      "Epoch 36:\n",
      "Train Loss: 68.1551, train Acc: 85.00%\n",
      "test Loss: 13.6366, test Acc: 86.33%\n",
      "--------------------------------------------------\n",
      "Epoch 37:\n",
      "Train Loss: 67.9827, train Acc: 85.32%\n",
      "test Loss: 13.5766, test Acc: 86.75%\n",
      "--------------------------------------------------\n",
      "Epoch 38:\n",
      "Train Loss: 67.6898, train Acc: 85.60%\n",
      "test Loss: 13.5447, test Acc: 86.83%\n",
      "--------------------------------------------------\n",
      "Epoch 39:\n",
      "Train Loss: 67.3732, train Acc: 85.85%\n",
      "test Loss: 13.4836, test Acc: 86.92%\n",
      "--------------------------------------------------\n",
      "Epoch 40:\n",
      "Train Loss: 67.0756, train Acc: 86.25%\n",
      "test Loss: 13.4424, test Acc: 87.00%\n",
      "--------------------------------------------------\n",
      "Epoch 41:\n",
      "Train Loss: 66.8704, train Acc: 86.43%\n",
      "test Loss: 13.4001, test Acc: 87.67%\n",
      "--------------------------------------------------\n",
      "Epoch 42:\n",
      "Train Loss: 66.6622, train Acc: 86.63%\n",
      "test Loss: 13.3424, test Acc: 87.75%\n",
      "--------------------------------------------------\n",
      "Epoch 43:\n",
      "Train Loss: 66.2869, train Acc: 86.88%\n",
      "test Loss: 13.2782, test Acc: 87.67%\n",
      "--------------------------------------------------\n",
      "Epoch 44:\n",
      "Train Loss: 65.9860, train Acc: 87.08%\n",
      "test Loss: 13.2385, test Acc: 87.92%\n",
      "--------------------------------------------------\n",
      "Epoch 45:\n",
      "Train Loss: 65.8496, train Acc: 87.22%\n",
      "test Loss: 13.2128, test Acc: 88.00%\n",
      "--------------------------------------------------\n",
      "Epoch 46:\n",
      "Train Loss: 65.6387, train Acc: 87.35%\n",
      "test Loss: 13.1660, test Acc: 88.50%\n",
      "--------------------------------------------------\n",
      "Epoch 47:\n",
      "Train Loss: 65.4667, train Acc: 87.70%\n",
      "test Loss: 13.1637, test Acc: 88.67%\n",
      "--------------------------------------------------\n",
      "Epoch 48:\n",
      "Train Loss: 65.4870, train Acc: 87.65%\n",
      "test Loss: 13.1152, test Acc: 88.83%\n",
      "--------------------------------------------------\n",
      "Epoch 49:\n",
      "Train Loss: 65.1145, train Acc: 87.80%\n",
      "test Loss: 13.0137, test Acc: 89.08%\n",
      "--------------------------------------------------\n",
      "Epoch 50:\n",
      "Train Loss: 64.7349, train Acc: 88.20%\n",
      "test Loss: 12.9985, test Acc: 89.42%\n",
      "--------------------------------------------------\n",
      "Training time: 0.9359402656555176 seconds\n",
      "-0.7923765222231547\n"
     ]
    }
   ],
   "source": [
    "\n",
    "if __name__ == \"__main__\":\n",
    "    warnings.filterwarnings(\"ignore\")\n",
    "    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "    os.environ['CUDA_LAUNCH_BLOCKING'] = '1'\n",
    "    begin_time = time.time()\n",
    "    num_class = 3\n",
    "    num_epochs = 50\n",
    "    batch_size = 64\n",
    "    learning_rate = 0.001\n",
    "    train, test, clsTR, clsTE= read_data()\n",
    "    run_RNN_module(device, num_class, num_epochs, batch_size, learning_rate, train, test, clsTR, clsTE)\n",
    "    end_time=time.time()\n",
    "    print((begin_time-end_time)/60)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "885fdd1e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
